﻿using System.Linq;
using Roslyn.Compilers.Common;
using Roslyn.Compilers.CSharp;
using RoslynHelper;

namespace YieldReturnYieldReturnManyReturn
{
    internal class YieldReturnYieldReturnManyReturnRewriter : SyntaxRewriter
    {
        private ISemanticModel _semanticModel;
        private bool _replaceReturn = false;
        private MethodDeclarationSyntax _method;
        private SyntaxNode _yieldReturnMany;

        public YieldReturnYieldReturnManyReturnRewriter(ISemanticModel semanticModel)
        {
            _semanticModel = semanticModel;
        }

        public override SyntaxNode Visit(SyntaxNode node)
        {
            if (node == _yieldReturnMany)
            {
                _yieldReturnMany = null;
                return null;
            }
            return base.Visit(node);
        }

        protected override SyntaxNode VisitMethodDeclaration(MethodDeclarationSyntax node)
        {
            _replaceReturn = new YieldReturnAndreturnMethodVisitor().Visit(node) == true;
            _method = node;
            var value = base.VisitMethodDeclaration(node);
            _replaceReturn = false;
            return value;
        }

        protected override SyntaxNode VisitReturnStatement(ReturnStatementSyntax node)
        {
            if (!_replaceReturn)
                return node;
            string variableName = _semanticModel.GetVariableName("yieldReturnVarLoop");
            return Syntax.Block(
                statements: Syntax.List<StatementSyntax>(
                    Syntax.ForEachStatement(
                        type: ((GenericNameSyntax)_method.ReturnType).TypeArgumentList.Arguments.First(),
                        identifier: Syntax.ParseToken(variableName),
                        expression: node.ExpressionOpt,
                        statement: Syntax.YieldStatement(SyntaxKind.YieldReturnStatement,
                            expressionOpt: Syntax.IdentifierName(variableName))),
                     Syntax.YieldStatement(SyntaxKind.YieldBreakStatement))
                );
        }

        protected override SyntaxNode VisitYieldStatement(YieldStatementSyntax node)
        {
            IdentifierNameSyntax identifier = node.ExpressionOpt as IdentifierNameSyntax;
            if (identifier != null && identifier.PlainName == "many")
            {
                var expStatement = node.Parent.ChildNodes().SkipWhile(n => n != node).Skip(1).FirstOrDefault() as ExpressionStatementSyntax;
                if (expStatement != null && !node.ToString().TrimEnd().EndsWith(";"))
                {
                    var exp = expStatement.Expression;
                    _yieldReturnMany = expStatement;
                    string variableName = _semanticModel.GetVariableName("yieldReturnVarLoop");
                    return Syntax.ForEachStatement(
                        type: ((GenericNameSyntax)_method.ReturnType).TypeArgumentList.Arguments.First(),
                        identifier: Syntax.ParseToken(variableName),
                        expression: exp,
                        statement: Syntax.YieldStatement(SyntaxKind.YieldReturnStatement,
                            expressionOpt: Syntax.IdentifierName(variableName))
                        );
                }
            }
            return node;
        }
    }
}
